import time
import numpy as np
import tensorflow as tf

with open("anna.txt", "r") as f:
    text = f.read()

vocabulary = set(text)

vocabulary_to_int = {c: i for i, c in enumerate(vocabulary)}
int_to_vocabulary = dict(enumerate(vocabulary))

encoded = np.array([vocabulary_to_int[w] for w in text], dtype = np.int32)


def get_batches(arr, num_seqs_for_single_batch, num_steps_for_single_seq):
    '''
    :param arr:
    :param num_seqs_for_single_batch:一个batch里的seq数量
    :param num_steps_for_single_seq:即sequence length
    :return:
    '''
    batch_word_size = num_seqs_for_single_batch * num_steps_for_single_seq          #单个batch的大小
    n_batches = int(len(arr) / batch_word_size)                                     #取整：忽略了每个batch最后不满足num_steps长度的内容
    arr = arr[: batch_word_size * n_batches]                                        #见上
    arr = arr.reshape((num_seqs_for_single_batch, -1))
    for n in range(0, arr.shape[1], num_steps_for_single_seq):
        x = arr[:, n:n+num_steps_for_single_seq]
        y = np.zeros_like(x)
        y[:, :-1], y[:, -1] = x[:, 1:], x[:, 0]                     #y：x的首列pop到最后一列
        yield x, y


def build_inputs(num_seqs_for_single_batch, num_steps_for_single_seq):
    inputs = tf.placeholder(tf.int32, shape=(num_seqs_for_single_batch, num_steps_for_single_seq), name="inputs")
    targets = tf.placeholder(tf.int32, shape=(num_seqs_for_single_batch, num_steps_for_single_seq), name="targets")

    #加入keep_prob：用来在后面控制dropout的保留结点数（去除不必要的网络节点达到正则化，减轻过拟合问题）
    keep_prob = tf.placeholder(tf.float32, name="keep_prob")

    return inputs, targets, keep_prob


def build_lstm(lstm_size, num_layers, num_seqs_for_single_batch, keep_prob):
    lstm = tf.nn.rnn_cell.BasicLSTMCell(lstm_size)                              #tf.contrib.rnn...
    drop = tf.nn.rnn_cell.DropoutWrapper(lstm, output_keep_prob=keep_prob)

    cell = tf.nn.rnn_cell.MultiRNNCell([drop for _ in range(num_layers)])
    initial_size = cell.zero_state(num_seqs_for_single_batch, tf.float32)

    return cell, initial_size


def build_output(lstm_output, in_size, out_size):
    seq_output = tf.concat(lstm_output, 1)                  #按axis=1(列)进行concat：[[1,2,3],[4,5,6]] -> [[1,2,3,4,5,6]]
    x = tf.reshape(seq_output, [-1, in_size])

    #连接LSTM输入到softmax layer
    with tf.variable_scope("softmax"):
        softmax_w = tf.Variable(tf.truncated_normal([in_size, out_size], stddev=0.1))
        softmax_b = tf.Variable(tf.zeros(out_size))

    logits = tf.matmul(x, softmax_w) + softmax_b

    out = tf.nn.softmax(logits, name="predictions")

    return out, logits


def build_loss(logits, targets, lstm_size, num_classes):
    y_one_hot = tf.one_hot(targets, num_classes)
    y_reshaped = tf.reshape(y_one_hot, logits.get_shape())

    loss = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=y_reshaped)
    loss = tf.reduce_mean(loss)

    return loss


def build_optimizer(loss, learning_rate, grad_clip):
    tvars = tf.trainable_variables()
    grads, _ = tf.clip_by_global_norm(tf.gradients(loss, tvars), grad_clip)
    train_op = tf.train.AdamOptimizer(learning_rate=learning_rate)
    optimizer = train_op.apply_gradients(zip(grads, tvars))

    return optimizer


class CharRNN:
    def __init__(self, num_classes, num_seq_for_signle_batch=64, num_steps_for_single_seq=50, lstm_size=128,
                 num_layers=2, learning_rate=0.001, grad_clip=5, sampling=False):
        if sampling:
            num_seq_for_signle_batch, num_steps_for_single_seq = 1, 1                    #采用SGD
        else:
            num_seq_for_signle_batch, num_steps_for_single_seq = num_seq_for_signle_batch, num_steps_for_single_seq

        tf.reset_default_graph()

        #输入层
        self.inputs, self.targets, self.keep_prob = build_inputs(num_seqs_for_single_batch=num_seq_for_signle_batch,
                                                                 num_steps_for_single_seq=num_steps_for_single_seq)

        #LSTM层
        cell, self.initial_state = build_lstm(lstm_size=lstm_size, num_layers=num_layers,
                                              num_seqs_for_single_batch=num_seq_for_signle_batch, keep_prob=self.keep_prob)

        #对输入进行one-hot编码
        x_one_hot = tf.one_hot(self.inputs, num_classes)

        #执行RNN
        outputs, state = tf.nn.dynamic_rnn(cell=cell, inputs=x_one_hot, initial_state=self.initial_state)
        self.final_state = state

        #预测结果
        self.prediction, self.logits = build_output(lstm_output=outputs, in_size=lstm_size, out_size=num_classes)

        #Loss和optimizer(with gradient clipping)
        self.loss = build_loss(logits = self.logits, targets=self.targets, lstm_size=lstm_size, num_classes=num_classes)
        self.optimizer = build_optimizer(loss=self.loss, learning_rate=learning_rate, grad_clip=grad_clip)


num_seq_for_signle_batch = 100
num_steps_for_single_seq = 100
lstm_size = 512
num_layers = 2
learning_rate = 0.001
keep_prob = 0.5

epochs = 20
save_every_n = 200

model = CharRNN(num_classes=len(vocabulary), num_seq_for_signle_batch=num_seq_for_signle_batch,
                num_steps_for_single_seq=num_steps_for_single_seq, lstm_size=lstm_size, num_layers=num_layers,
                learning_rate=learning_rate)

saver = tf.train.Saver(max_to_keep=100)
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    counter = 0
    for e in range(epochs):
        #Train network
        new_state = sess.run(model.initial_state)
        loss = 0
        for x, y in get_batches(encoded, num_seq_for_signle_batch, num_steps_for_single_seq):
            counter += 1
            start = time.time()
            feed = {model.inputs: x, model.targets: y, model.keep_prob: keep_prob, model.initial_state: new_state}
            batch_loss, new_state, _ = sess.run([model.loss, model.final_state, model.optimizer], feed_dict=feed)
            end = time.time()

            if counter % 50 == 0:
                print('轮数: {}/{}... '.format(e + 1, epochs),
                      '训练步数: {}... '.format(counter),
                      '训练误差: {:.4f}... '.format(batch_loss),
                      '{:.4f} sec/batch'.format((end - start)))

            if counter % save_every_n == 0:
                saver.save(sess, "checkpoints/i{}_l{}.ckpt".format(counter, lstm_size))
    saver.save(sess, "checkpoints/i{}_l{}.ckpt".format(counter, lstm_size))
